import torch
from torch import nn
from torch.autograd import Variable
from torch.utils.data import DataLoader
import torchvision.datasets as dataset
import torchvision.transforms as transforms

batch_size = 100
# 加载数据集
train_dataset = dataset.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = dataset.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=True)
#  封装数据集
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)

# 初始化超参数
input_size = 784
hidden_size = 500
num_classes = 10  # 类别数目


# 定义神经网络模型
class module(nn.Module):
    def __init__(self, input_size, hidden_size, output_num):
        super(module, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_num)

    # 前向传播
    def forward(self, x):
        out = self.fc1(x)
        out = torch.relu(out)
        out = self.fc2(out)
        return out

module = module(input_size, hidden_size, num_classes)

# 超参数初始化
lr = 1e-1
epochs = 5
# 定义损失函数及优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(module.parameters(), lr=lr)
# 模型开始训练
for epoch in range(epochs + 1):
    print("==============第 {} 轮 训练开始==============".format(epoch + 1))
    for i, (images, labels) in enumerate(train_loader):
        images = Variable(images.view(-1, 28 * 28))
        labels = Variable(labels)
        outputs = module(images)
        # 计算损失
        loss = criterion(outputs, labels)
        # 梯度清0
        optimizer.zero_grad()
        # 反向传播
        loss.backward()
        # 参数更新
        optimizer.step()  # update parameters
        #
        if i % 100 == 0:
            print("交叉熵损失为: %.5f" % loss.item())

# 利用训练好的模型进行预测
T = 0  # 总共测试集样本个数
CC = 0  # 测试集正确识别的个数
for images, labels in test_loader:
    images = Variable(images.view(-1, 28 * 28))
    labels = Variable(labels)
    outputs = module(images)
    _, predicts = torch.max(outputs.data, 1)
    T += labels.size(0)
    CC += (predicts == labels).sum()

print("识别准确率为：%.2f %%" % (100 * CC / T))


